'''
Author: 渔火Arcadia  https://github.com/yhArcadia
Date: 2025-04-18 00:13:00
LastEditors: 渔火Arcadia
LastEditTime: 2025-04-18 22:53:20
FilePath: /example/train.py
Description: 

Copyright (c) 2025 by 渔火Arcadia 1761869682@qq.com, All Rights Reserved. 
'''
from model import mymodel
import torch
import torch.nn as nn
from utils.loader import train_loader

# 训练轮次
epochs = 20

# 加载模型、优化器和损失函数
model = mymodel()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001) # Adam优化器，用于更新模型参数。
loss_func = nn.CrossEntropyLoss() # 交叉熵损失函数，适用于多分类问题，其输入是模型输出的概率分布和真实标签。

# 加载图像，设置真值标签
imgs = train_loader()
truth = torch.tensor([0,1,2,2,3])

 # 模型设置为训练模式, 这样会启用Dropout和BatchNorm等训练时的特殊行为，如随机丢弃神经元和计算均值和方差。
model.train()

# 训练模型
for epoch in range(epochs):
    # 每一轮，图片输入到模型得到损失值
    res = model(imgs) # 将图片喂给模型，得到预测值（也就是对于4个分类对应的概率分布）
    loss = loss_func(res,truth) #预测值和真值标签对比，得到损失值

    # 反向传播三件套
    optimizer.zero_grad()  # 1. 清空梯度
    loss.backward()        # 2. 反向传播计算梯度
    optimizer.step()       # 3. 更新参数

    print(f'\n第{epoch+1}次训练时，输入的5张图片各自对于4个分类的预测值分布：\n', res)

# 最后，保存模型文件
torch.save(model, './model.pth')
print("训练完毕")